package com.liuyun.asymmetric.ecdh;

import lombok.SneakyThrows;
import org.bouncycastle.jce.provider.BouncyCastleProvider;

import javax.crypto.Cipher;
import javax.crypto.KeyAgreement;
import javax.crypto.SecretKey;
import javax.crypto.spec.SecretKeySpec;
import java.security.Key;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.Security;
import java.security.interfaces.ECPrivateKey;
import java.security.interfaces.ECPublicKey;
import java.security.spec.ECParameterSpec;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.HashMap;
import java.util.Map;

/**
 * 非对称算法DH：用于密钥交换，解决对称解密算法管理密钥的难题
 * <p>
 * Copyright: (C), 2021/1/24 22:04
 * <p>
 * <p>
 * Company Information & Data Security Solutions Co., Ltd.
 *
 * @author LiuYun liujw@idss-cn.com
 * @version 1.0
 */
public class EcDhUtil {
    static {
        // JDK 不支持EC和ECDH
        Security.addProvider(new BouncyCastleProvider());
    }

    public static final String PUBLIC_KEY = "publicKey";
    public static final String PRIVATE_KEY = "privateKey";
    public static final String KEY_ALGORITHM = "ECDH";
    // 本地密钥算法（对称加密密钥算法），Blowfish、RC2、RC4等
    public static final String SECRET_ALGORITHM = "Blowfish";
    // ECDH算法长度默认256，可取范围：112-571
    public static final int KEY_SIZE = 256;

    /**
     * 初始化甲方密钥对
     *
     * @return
     */
    @SneakyThrows
    public static Map<String, Key> initKey() {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(KEY_ALGORITHM,"BC");
        keyPairGenerator.initialize(KEY_SIZE);
        KeyPair keyPair = keyPairGenerator.generateKeyPair();
        ECPrivateKey privateKey = (ECPrivateKey) keyPair.getPrivate();
        ECPublicKey publicKey = (ECPublicKey) keyPair.getPublic();
        Map<String, Key> keyMap = new HashMap<>(2);
        keyMap.put(PUBLIC_KEY, publicKey);
        keyMap.put(PRIVATE_KEY, privateKey);
        return keyMap;
    }

    /**
     * 初始化乙方密钥对：需要传入甲方的公钥
     *
     * @return
     */
    @SneakyThrows
    public static Map<String, Key> initKey(byte[] key) {
        // 将甲方key转换为甲方的公钥对象
        X509EncodedKeySpec x509EncodedKeySpec = new X509EncodedKeySpec(key);
        KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM,"BC");
        ECPublicKey pubKey = (ECPublicKey) keyFactory.generatePublic(x509EncodedKeySpec);
        // 由甲方公钥构建乙方密钥
        ECParameterSpec params = pubKey.getParams();

        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(KEY_ALGORITHM,"BC");
        keyPairGenerator.initialize(params);
        KeyPair keyPair = keyPairGenerator.generateKeyPair();
        ECPrivateKey privateKey = (ECPrivateKey) keyPair.getPrivate();
        ECPublicKey publicKey = (ECPublicKey) keyPair.getPublic();
        Map<String, Key> keyMap = new HashMap<>(2);
        keyMap.put(PUBLIC_KEY, publicKey);
        keyMap.put(PRIVATE_KEY, privateKey);
        return keyMap;
    }

    /**
     * 对称加密数据
     *
     * @param data
     * @param key
     * @return
     */
    @SneakyThrows
    public static byte[] encrypt(byte[] data, byte[] key) {
        // 生成本地密钥
        SecretKey secretKey = new SecretKeySpec(key, SECRET_ALGORITHM);
        // 加密
        Cipher cipher = Cipher.getInstance(secretKey.getAlgorithm());
        cipher.init(Cipher.ENCRYPT_MODE, secretKey);
        return cipher.doFinal(data);
    }

    /**
     * 对称解密数据
     *
     * @param data
     * @param key
     * @return
     */
    @SneakyThrows
    public static byte[] decrypt(byte[] data, byte[] key) {
        // 生成本地密钥
        SecretKey secretKey = new SecretKeySpec(key, SECRET_ALGORITHM);
        // 加密
        Cipher cipher = Cipher.getInstance(secretKey.getAlgorithm());
        cipher.init(Cipher.DECRYPT_MODE, secretKey);
        return cipher.doFinal(data);
    }

    @SneakyThrows
    public static byte[] buildSecretKey(byte[] publicKey, byte[] privateKey) {
        KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM,"BC");
        // 初始化对端公钥
        X509EncodedKeySpec x509EncodedKeySpec = new X509EncodedKeySpec(publicKey);
        PublicKey pubKey = keyFactory.generatePublic(x509EncodedKeySpec);

        // 初始化本地私钥
        PKCS8EncodedKeySpec pkcs8EncodedKeySpec = new PKCS8EncodedKeySpec(privateKey);
        PrivateKey priKey = keyFactory.generatePrivate(pkcs8EncodedKeySpec);

        // 使用对端公钥和本地私钥，协商生成本地密钥（对称算法的密钥）
        KeyAgreement keyAgreement = KeyAgreement.getInstance(keyFactory.getAlgorithm(),"BC");
        keyAgreement.init(priKey);
        keyAgreement.doPhase(pubKey, true);
        SecretKey secretKey = keyAgreement.generateSecret(SECRET_ALGORITHM);
        return secretKey.getEncoded();
    }

    public static byte[] getPublicKey(Map<String, Key> keyMap) {
        return keyMap.get(PUBLIC_KEY).getEncoded();
    }

    public static byte[] getPrivateKey(Map<String, Key> keyMap) {
        return keyMap.get(PRIVATE_KEY).getEncoded();
    }
}
